{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZPrsyJiUBi6M"
      },
      "source": [
        "Licensed under the Apache License, Version 2.0"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qZF3Ypk7dTwq"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "if 'google.colab' in sys.modules:\n",
        "  from google.colab import auth\n",
        "  auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "y9RJIjajUxcC"
      },
      "outputs": [],
      "source": [
        "def parse_example(serialized_example: bytes) -\u003e dict[str, tf.Tensor]:\n",
        "  # GOES-16 ABI band 8 to 16\n",
        "  goes_bands = ('upper_vapor', 'mid_vapor', 'low_vapor', 'cloud_top', 'ozone',\n",
        "                'data_10um', 'data_11um', 'data_12um', 'co2')\n",
        "  feature_spec = {band: tf.io.FixedLenFeature([], tf.string) for band in goes_bands}\n",
        "  feature_spec.update({\n",
        "      'brightness_temperature_difference':  tf.io.FixedLenFeature([], tf.string),  # data_11um - data_12um\n",
        "      'human_individual_masks': tf.io.FixedLenFeature([], tf.string),  # individual labeled masks\n",
        "      'human_pixel_masks': tf.io.FixedLenFeature([], tf.string),  # majority vote from human_indvidual_masks      \n",
        "      'n_times_before': tf.io.FixedLenFeature([], tf.int64),\n",
        "      'n_times_after': tf.io.FixedLenFeature([], tf.int64),\n",
        "      # Projection params\n",
        "      'projection_wkt': tf.io.FixedLenFeature([], tf.string),\n",
        "      'col_min': tf.io.FixedLenFeature([], tf.float32),\n",
        "      'row_min': tf.io.FixedLenFeature([], tf.float32),\n",
        "      'col_size': tf.io.FixedLenFeature([], tf.float32),\n",
        "      'row_size': tf.io.FixedLenFeature([], tf.float32),\n",
        "      # Timestamp\n",
        "      'timestamp': tf.io.FixedLenFeature([], tf.int64),  # approximate timestamp\n",
        "      'satellite_scan_starts': tf.io.VarLenFeature(tf.int64),  # timestamp from original file\n",
        "  })\n",
        "  features = tf.io.parse_single_example(serialized_example, feature_spec)\n",
        "  for key in [*goes_bands, 'brightness_temperature_difference']:\n",
        "    features[key] = tf.io.parse_tensor(features[key], tf.double)\n",
        "  features['human_pixel_masks'] = tf.io.parse_tensor(features['human_pixel_masks'], tf.int32)\n",
        "  features['human_individual_masks'] = tf.io.parse_tensor(features['human_individual_masks'], tf.int32)\n",
        "  return features"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QpVZmmNIG3Nk"
      },
      "source": [
        "## False color image\n",
        "In order to view contrails in GOES, we use the \"ash\" color scheme. This color scheme was originally developed for viewing volcanic ash in the atmosphere but is also useful for viewing thin cirrus, including contrails. In this color scheme, contrails appear in the image as dark blue.\n",
        "\n",
        "Note that we use a modified version of the ash color scheme here, developed by Kulik et al., which uses slightly different bands and bounds tuned for contrails.\n",
        "\n",
        "References:\n",
        "\n",
        "Ash Color Scheme (page 7): https://eumetrain.org/sites/default/files/2020-05/RGB_recipes.pdf"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WGzd_H80GhP6"
      },
      "outputs": [],
      "source": [
        "_T11_BOUNDS = (243, 303)\n",
        "_CLOUD_TOP_TDIFF_BOUNDS = (-4, 5)\n",
        "_TDIFF_BOUNDS = (-4, 2)\n",
        "\n",
        "def normalize_range(data, bounds):\n",
        "  \"\"\"Maps data to the range [0, 1].\"\"\"\n",
        "  return (data - bounds[0]) / (bounds[1] - bounds[0])\n",
        "\n",
        "def false_color_image(brightness_temperatures):\n",
        "  \"\"\"Generates ash false color image from GOES brightness temperatures.\"\"\"\n",
        "  r = normalize_range(brightness_temperatures['data_12um'] - brightness_temperatures['data_11um'], _TDIFF_BOUNDS)\n",
        "  g = normalize_range(brightness_temperatures['data_11um'] - brightness_temperatures['cloud_top'], _CLOUD_TOP_TDIFF_BOUNDS)\n",
        "  b = normalize_range(brightness_temperatures['data_11um'], _T11_BOUNDS)\n",
        "  return np.clip(np.stack([r, g, b], axis=-1), 0, 1)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eSpmD22RdVF6"
      },
      "outputs": [],
      "source": [
        "dataset = tf.data.TFRecordDataset(tf.io.gfile.glob('gs://goes_contrails_dataset/20230419/tfrecords/train.tfrecords-00000-of-00100'))\n",
        "dataset = dataset.map(parse_example)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "reSh5o_5dXYc"
      },
      "outputs": [],
      "source": [
        "count = 0\n",
        "for features in dataset.as_numpy_iterator():\n",
        "  # Skip examples without contrails for visualization\n",
        "  if features['human_pixel_masks'].sum() == 0:\n",
        "    continue\n",
        "  n_times_before = features['n_times_before']\n",
        "  false_color_images = false_color_image(features)\n",
        "  plt.figure(figsize=(12, 6))\n",
        "  plt.subplot(1, 2, 1)\n",
        "  plt.imshow(false_color_images[..., n_times_before, :])\n",
        "  plt.subplot(1, 2, 2)\n",
        "  plt.imshow(features['human_pixel_masks'].squeeze(-1))\n",
        "  plt.show()\n",
        "  count += 1\n",
        "  if count \u003e= 10:\n",
        "    break"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "last_runtime": {
        "build_target": "//research/energy/contrails/colab:kernel",
        "kind": "private"
      },
      "private_outputs": true,
      "provenance": [
        {
          "file_id": "1JPUZgPTjY6SbNkBSbezhx8BYX9EitAcA",
          "timestamp": 1682745713184
        }
      ]
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
